//! Pseudo-SSA form for lowering expressions to linear code.

use std::collections::HashMap;

use crate::mir::{EnvExpr, TaggedExpr, TypedExpr};
use crate::tysys::{
    ArgRef, BinOp, FieldRef, FnContext, IntEncoding, LiteralValue, LocalRef, MonOp, ValueType,
};

/// Placeholder alias meant to make code less hard to rewrite.
pub type VarRef = LocalRef;

impl From<EnvExpr> for SsaArg {
    fn from(value: EnvExpr) -> Self {
        match value {
            EnvExpr::Field(i) => Self::Field(i),
            EnvExpr::Arg(i) => Self::Arg(i),
            EnvExpr::Local(i) => Self::Local(i),
        }
    }
}

/// SSA expression, not a AST expression.
#[derive(Clone, Debug)]
pub enum SsaExpr {
    /// Stores a literal.
    Literal(LiteralValue),

    /// Copy the value from another assignment to a new one.
    Copy(SsaArg),

    /// Stores the result of a mono-op.
    Monop(MonOp, SsaArg),

    /// Stores the result of a bin-op.
    Binop(BinOp, SsaArg, SsaArg),

    /// Depending on a condition arg, evals either of the two blocks and stores their result.
    Cond(SsaArg, Box<Block>, Box<Block>),

    /// Simple ternary expression, used when a full block is unnecessary like
    /// with picking between two consts/args.
    Ternary(SsaArg, SsaArg, SsaArg),

    /// Casts an argument to a type, whatever the native representation in the
    /// target is.
    Cast(ValueType, SsaArg),

    /// Evaluates to a buffer's length.
    BufferLen(SsaArg),

    /// Stores a checked read from a buffer.
    ///
    /// (encoding, buf, offset)
    ReadInt(IntEncoding, SsaArg, SsaArg),

    /// Stores a slice checked read from a buffer.
    ///
    /// (buffer, offset, len)
    ReadBuf(SsaArg, SsaArg, SsaArg),
}

/// An entry within an SSA block.
#[derive(Clone, Debug)]
pub enum SsaEntry {
    /// Assignment to a new variable.
    // TODO replace with AssignmentTarget
    Assignment(AssignmentTarget, SsaExpr),

    /// Write at some value to an offset in a buffer.
    ///
    /// (buf, offset, value)
    WriteBuf(AssignmentTarget, SsaArg, SsaArg),

    /// Abort if this if this value is truthy.
    AbortIf(SsaArg, AbortReason),

    /// Returns from the function with no value, probably breaking other
    /// expectations.
    Return,

    /// Returns from the function with a value immediately, probably breaking
    /// other expectations.
    ReturnValue(SsaArg),
}

/// A possible target of an assignment, irrespective of indexing or whatever.
#[derive(Clone, Debug)]
pub enum AssignmentTarget {
    /// A field in a struct, if in a struct.
    Field(FieldRef),

    /// A local variable.
    Local(VarRef),
}

/// A block that gets executed linearly.
#[derive(Clone, Debug)]
pub struct Block {
    entries: Vec<SsaEntry>,
    result: Option<SsaArg>,
}

impl Block {
    pub fn entries(&self) -> &[SsaEntry] {
        &self.entries
    }

    pub fn result(&self) -> Option<&SsaArg> {
        self.result.as_ref()
    }
}

pub struct ThunkContext {
    var_tys: Vec<ValueType>,
}

/// Environment context that args might be resolved in.
#[derive(Clone, Debug)]
pub struct TypeContextData {
    /// Vec of argument types.
    arg_tys: Vec<ValueType>,

    /// Var types as configured.
    var_tys: Vec<ValueType>,
}

impl TypeContextData {
    pub fn new(arg_tys: Vec<ValueType>, var_tys: Vec<ValueType>) -> Self {
        Self { arg_tys, var_tys }
    }

    pub fn new_empty() -> Self {
        Self::new(Vec::new(), Vec::new())
    }

    pub fn num_args(&self) -> u16 {
        self.arg_tys.len() as u16
    }

    pub fn get_arg_ty(&self, idx: ArgRef) -> Option<&ValueType> {
        self.arg_tys.get(idx.inner() as usize)
    }

    pub fn get_var_ty(&self, vr: LocalRef) -> Option<&ValueType> {
        self.var_tys.get(vr.inner() as usize)
    }
}

/// The bulk of function as a series of SSA steps with types.
#[derive(Clone, Debug)]
pub struct Function {
    /// The context for the function.
    ctx: TypeContextData,

    /// Return type.
    return_ty: ValueType,

    /// If the function can mutate the backing buffer.
    muts_buf: bool,

    /// Root SSA block.
    root_block: Block,
}

impl Function {
    pub fn arg_num(&self) -> u16 {
        self.ctx.num_args()
    }

    pub fn return_ty(&self) -> &ValueType {
        &self.return_ty
    }

    pub fn is_mutating(&self) -> bool {
        self.muts_buf
    }

    pub fn root_block(&self) -> &Block {
        &self.root_block
    }

    pub fn get_arg_ty(&self, idx: ArgRef) -> Option<&ValueType> {
        self.ctx.get_arg_ty(idx)
    }

    pub fn get_var_ty(&self, vr: LocalRef) -> Option<&ValueType> {
        self.ctx.get_var_ty(vr)
    }
}

impl FnContext for Function {
    fn is_container(&self) -> bool {
        // This *is* a container *here*.
        true
    }

    fn get_arg_type(&self, arg: ArgRef) -> Option<ValueType> {
        self.get_arg_ty(arg).cloned()
    }

    fn get_local_type(&self, local: LocalRef) -> Option<ValueType> {
        self.get_var_ty(local).cloned()
    }

    fn get_instance_field_type(&self, field: FieldRef) -> Option<ValueType> {
        // FIXME this is hardcoded for buf wrapper types
        if field.inner() == 0 {
            // FIXME this ignores what we expect the length to be
            Some(ValueType::ByteSlice)
        } else {
            None
        }
    }

    fn num_args(&self) -> u16 {
        self.ctx.arg_tys.len() as u16
    }

    fn num_instance_fields(&self) -> u16 {
        // This is always just the buffer for now.
        1
    }
}

/// Maintains state local to the current block we're building.
struct BlockBuilder {
    entries: Vec<SsaEntry>,
}

impl BlockBuilder {
    fn new() -> Self {
        Self {
            entries: Vec::new(),
        }
    }

    fn finalize(self, result: Option<SsaArg>) -> Block {
        Block {
            entries: self.entries,
            result,
        }
    }
}

pub struct SsaBuilder {
    /// Maximum argument seen in any expressions.
    seen_args: u16,

    /// Types of arguments, as attested.
    arg_types: HashMap<ArgRef, ValueType>,

    /// Next variable for locals.
    next_var: u16,

    /// Types of locals, as attested.
    var_types: HashMap<VarRef, ValueType>,

    /// Return type that we were told about.
    return_ty: Option<ValueType>,

    /// If the function mutates the backing buffer.
    muts_buf: bool,

    /// Stack of block contexts we're constructing.  This will always have the
    /// root block at index 0 and contain only the root block when finalizing.
    blocks: Vec<BlockBuilder>,
}

impl SsaBuilder {
    pub fn new() -> Self {
        Self {
            seen_args: 0,
            arg_types: HashMap::new(),
            next_var: 0,
            var_types: HashMap::new(),
            return_ty: None,
            muts_buf: false,
            blocks: vec![BlockBuilder::new()],
        }
    }

    fn witness_arg(&mut self, arg_idx: ArgRef) {
        self.seen_args = u16::max(self.seen_args, arg_idx.inner() + 1)
    }

    pub fn set_arg_type(&mut self, arg_idx: ArgRef, ty: ValueType) {
        self.witness_arg(arg_idx);

        if let Some(aty) = self.arg_types.get(&arg_idx) {
            if *aty != ty {
                // TODO convert to normal
                panic!("ssa: arg {arg_idx:?} stored as {aty:?} but just used as {ty:?}");
            }
        }

        self.arg_types.insert(arg_idx, ty);
    }

    pub fn set_return_type(&mut self, ty: ValueType) {
        if let Some(rty) = &self.return_ty {
            if rty != &ty {
                panic!("ssa: return type stored as {rty:?} but just tried to use as {ty:?}");
            }
        }
        self.return_ty = Some(ty);
    }

    fn set_mut_flag(&mut self) {
        self.muts_buf = true;
    }

    pub fn get_var(&mut self) -> LocalRef {
        let v = self.next_var;
        self.next_var += 1;
        v.into()
    }

    fn cur_block_mut(&mut self) -> &mut BlockBuilder {
        self.blocks.last_mut().expect("ssabuilder: bottomed out")
    }

    pub fn put_entry(&mut self, entry: SsaEntry) {
        // If we ever see this entry then we know we mutate the buffer.
        if let SsaEntry::WriteBuf(_, _, _) = &entry {
            self.set_mut_flag();
        }

        let cb = self.cur_block_mut();
        cb.entries.push(entry);
    }

    pub fn put_expr_assign(&mut self, expr: SsaExpr, ty: ValueType) -> VarRef {
        let v = self.get_var();
        let vat = AssignmentTarget::Local(v);
        self.put_entry(SsaEntry::Assignment(vat, expr));
        self.var_types.insert(v, ty);
        v
    }

    pub fn push_block_ctx(&mut self) {
        self.blocks.push(BlockBuilder::new());
    }

    pub fn pop_block_ctx(&mut self, result: Option<SsaArg>) -> Block {
        let block = self.blocks.pop().expect("ssabuilder: bottomed out");
        block.finalize(result)
    }

    pub fn get_arg_tys(&self) -> Vec<ValueType> {
        (0..self.seen_args)
            .map(|i| {
                let arg = ArgRef::from(i);
                let ty = self.arg_types.get(&arg).cloned();
                if ty.is_none() {
                    panic!("ssa: missing type assignment for arg {i}");
                }
                ty.clone().unwrap()
            })
            .collect::<Vec<_>>()
    }

    pub fn get_local_ty(&self, vr: VarRef) -> Option<&ValueType> {
        self.var_types.get(&vr)
    }

    pub fn get_local_tys(&self) -> Vec<ValueType> {
        (0..self.next_var)
            .map(|i| {
                let lr = LocalRef::from(i);
                let ty = self.var_types.get(&lr).cloned();
                if ty.is_none() {
                    panic!("ssa: missing type assignment for local var {i}");
                }
                ty.unwrap()
            })
            .collect()
    }

    pub fn finalize_function(mut self, result: Option<SsaArg>) -> Function {
        if self.blocks.len() != 1 {
            panic!(
                "ssa: invalid builder state, have {} remaining blocks",
                self.blocks.len()
            );
        }

        let root_block = self.pop_block_ctx(result);
        let ctx = TypeContextData::new(self.get_arg_tys(), self.get_local_tys());
        Function {
            ctx,
            return_ty: self.return_ty.unwrap_or(ValueType::Void),
            muts_buf: self.muts_buf,
            root_block,
        }
    }

    pub fn finalize_thunk(mut self) -> BufThunk {
        if self.blocks.len() != 1 {
            panic!(
                "ssa: invalid builder state, have {} remaining blocks",
                self.blocks.len()
            );
        }

        let root_block = self.pop_block_ctx(None);
        let ctx = ThunkContext {
            var_tys: self.get_local_tys(),
        };

        BufThunk::new(ctx, root_block)
    }
}

/// Converts an expr for querying a value to SSA form.
pub fn convert_expr_to_getter_fun(expr: &TaggedExpr<ValueType>) -> Function {
    let mut builder = SsaBuilder::new();
    let final_vr = flatten_expr_rec(expr, &mut builder);
    let final_vr_ty = builder.get_local_ty(final_vr).expect("ssa: result varref");
    builder.set_return_type(final_vr_ty.clone());
    builder.finalize_function(Some(SsaArg::Local(final_vr)))
}

/// Recursively emits the evalutation steps for an expression returning a var
/// ref in the SSA builder.
pub fn flatten_expr_rec(expr: &TaggedExpr<ValueType>, builder: &mut SsaBuilder) -> VarRef {
    match expr {
        // This branch isn't used much since we usually short circuit the arg evaluation.
        TaggedExpr::Const(ty, v) => builder.put_expr_assign(SsaExpr::Literal(*v), ty.clone()),

        TaggedExpr::Env(ty, env) => {
            if let EnvExpr::Arg(arg) = env {
                builder.set_arg_type(*arg, ty.clone());
            }

            builder.put_expr_assign(SsaExpr::Copy(SsaArg::from(*env)), ty.clone())
        }

        TaggedExpr::Monop(ty, op, e) => {
            let arg = conv_arg_rec(e, builder);
            builder.put_expr_assign(SsaExpr::Monop(*op, arg), ty.clone())
        }

        TaggedExpr::Binop(ty, op, le, re) => {
            let la = conv_arg_rec(le, builder);
            let ra = conv_arg_rec(re, builder);
            builder.put_expr_assign(SsaExpr::Binop(*op, la, ra), ty.clone())
        }

        TaggedExpr::Cond(ty, ce, te, ee) => {
            let ca = conv_arg_rec(ce, builder);

            builder.push_block_ctx();
            let ta = conv_arg_rec(te, builder);
            let tb = builder.pop_block_ctx(Some(ta));

            builder.push_block_ctx();
            let ea = conv_arg_rec(ee, builder);
            let eb = builder.pop_block_ctx(Some(ea));

            builder.put_expr_assign(SsaExpr::Cond(ca, Box::new(tb), Box::new(eb)), ty.clone())
        }

        TaggedExpr::Cast(ty, target_ty, e) => {
            let arg = conv_arg_rec(e, builder);
            builder.put_expr_assign(SsaExpr::Cast(target_ty.clone(), arg), ty.clone())
        }

        TaggedExpr::BufferLen(ty, ebuf) => {
            let arg = conv_arg_rec(ebuf, builder);
            builder.put_expr_assign(SsaExpr::BufferLen(arg), ty.clone())
        }

        TaggedExpr::BufferRead(ty, bufe, se, le) => {
            let ba = conv_arg_rec(bufe, builder);
            let sa = conv_arg_rec(se, builder);
            let la = conv_arg_rec(le, builder);
            builder.put_expr_assign(SsaExpr::ReadBuf(ba, sa, la), ty.clone())
        }

        TaggedExpr::BufferReadInt(ty, bufe, oe, ienc) => {
            let ba = conv_arg_rec(bufe, builder);
            let oa = conv_arg_rec(oe, builder);
            builder.put_expr_assign(SsaExpr::ReadInt(*ienc, ba, oa), ty.clone())
        }
    }
}

/// Tries to short-circuit evaluation to avoid creating another local if the
/// expression is a constant or a function argument.
fn conv_arg_rec(expr: &TypedExpr, builder: &mut SsaBuilder) -> SsaArg {
    match expr {
        TaggedExpr::Const(_, v) => SsaArg::Literal(v.clone()),
        TaggedExpr::Env(ty, e) => {
            if let EnvExpr::Arg(arg) = e {
                builder.set_arg_type(*arg, ty.clone());
            }

            SsaArg::from(*e)
        }

        e => SsaArg::Local(flatten_expr_rec(e, builder)),
    }
}

/// A thunk that operates over an implied buffer.
pub struct BufThunk {
    ctx: ThunkContext,
    root_block: Block,
}

impl BufThunk {
    pub fn new(ctx: ThunkContext, root_block: Block) -> Self {
        Self { ctx, root_block }
    }

    pub fn context(&self) -> &ThunkContext {
        &self.ctx
    }

    pub fn root_block(&self) -> &Block {
        &self.root_block
    }
}

#[cfg(test)]
mod tests {
    use crate::mir::Expr;
    use crate::test_util;
    use crate::tysys::{typecheck_expr, ArgRef, ContainerData, MethodData, ValueType};

    use super::convert_expr_to_getter_fun;

    #[test]
    fn test_ssab_simple() {
        let cont = ContainerData::new(vec![ValueType::ByteSlice]);
        let meth = MethodData::new(&cont, Vec::new(), Vec::new());

        let e = test_util::get_012_cond_expr(Expr::new_arg(ArgRef::from(0)));

        let ie = e.to_int_tag();
        eprintln!("{ie:#?}");

        let te = typecheck_expr(&ie, &meth).expect("test: typecheck");
        eprintln!("{te:#?}");

        let ssa = convert_expr_to_getter_fun(&te);
        eprintln!("{ssa:#?}");
        assert_eq!(ssa.return_ty, ValueType::U8, "unexpected result type");
    }
}
